# 用于DMA的画图，和task-num的曲线这种。

import csv
import matplotlib.pyplot as plt
import os

# 从CSV文件中读取数据
def read_csv_data(file_path):
    data = []
    with open(file_path, newline='') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            data.append({
                'dma_task_num': int(row['dma_task_num']),
                'dma_task_size': int(row['dma_task_size']),
                'Speed_Mbps': float(row['Speed_Mbps'])
            })
    return data

# 绘制传输速度随dma_task_sizes变化的曲线
def plot_speed_vs_task_size(data, fixed_task_nums, save_dir):
    plt.figure(figsize=(10, 6))
    
    for task_num in fixed_task_nums:
        sizes = []
        speeds = []
        for entry in data:
            if entry['dma_task_num'] == task_num:
                sizes.append(entry['dma_task_size'])
                speeds.append(entry['Speed_Mbps'])
        
        plt.plot(sizes, speeds, marker='o', label=f'dma_task_num={task_num}')
    
    plt.xscale('log', base=2)  # 使用对数刻度以显示所有2的幂次方
    plt.xticks(sizes, [str(size) for size in sizes], rotation=45)
    plt.xlabel('dma_task_size')
    plt.ylabel('Speed (Mbps)')
    plt.title('Speed vs dma_task_size for different dma_task_nums')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    
    # 保存图像
    plot_path = os.path.join(save_dir, 'speed_vs_task_size.png')
    plt.savefig(plot_path)
    plt.close()

# 绘制传输速度随dma_task_nums变化的曲线
def plot_speed_vs_task_num(data, fixed_task_sizes, save_dir):
    plt.figure(figsize=(10, 6))
    
    for task_size in fixed_task_sizes:
        nums = []
        speeds = []
        for entry in data:
            if entry['dma_task_size'] == task_size:
                nums.append(entry['dma_task_num'])
                speeds.append(entry['Speed_Mbps'])
        
        plt.plot(nums, speeds, marker='o', label=f'dma_task_size={task_size}')
    
    plt.xscale('log', base=2)  # 使用对数刻度以显示所有2的幂次方
    plt.xticks(nums, [str(num) for num in nums], rotation=45)
    plt.xlabel('dma_task_num')
    plt.ylabel('Speed (Mbps)')
    plt.title('Speed vs dma_task_num for different dma_task_sizes')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    
    # 保存图像
    plot_path = os.path.join(save_dir, 'speed_vs_task_num.png')
    plt.savefig(plot_path)
    plt.close()

# 读取数据
data = read_csv_data('../results_raw/transfer_rates.csv')

# 固定的dma_task_nums和dma_task_sizes
fixed_task_nums = [1, 2, 4, 8, 16, 32, 64]  # 根据需要选择几个固定的任务数
fixed_task_sizes = [8, 1024, 65536, 131072, 262144, 524288, 1048576, 2097152]  # 根据需要选择几个固定的任务大小

# 指定保存图像的目录
save_dir = '../plots/dma/'
os.makedirs(save_dir, exist_ok=True)

# 绘制图形并保存
plot_speed_vs_task_size(data, fixed_task_nums, save_dir)
plot_speed_vs_task_num(data, fixed_task_sizes, save_dir)
